import { BaseAgent, BaseAgentConfig, AgentCommandConfig } from "./base";
import { ModelConfig } from "./utils";
import {
  AgentType,
  GrokConfig,
  GrokResponse,
  GrokStreamCallbacks,
  Conversation,
  ModelProvider,
} from "../types";

export class GrokAgent extends BaseAgent {
  private grokApiKey: string;
  private model?: string;
  private baseUrl?: string;

  private escapePrompt(prompt: string): string {
    // Comprehensive escaping for bash double quotes - production ready
    return prompt
      .replace(/\\/g, '\\\\')    // Escape backslashes FIRST
      .replace(/"/g, '\\"')      // Escape double quotes
      .replace(/'/g, "'\\''")    // Escape single quotes (close quote, escaped quote, reopen)
      .replace(/\$/g, '\\$')     // Escape dollar signs
      .replace(/`/g, '\\`')      // Escape backticks (command substitution)
      .replace(/!/g, '\\!')      // Escape exclamation (history expansion)
      .replace(/\n/g, '\\n')     // Escape newlines
      .replace(/\r/g, '\\r')     // Escape carriage returns
      .replace(/\t/g, '\\t');    // Escape tabs
  }

  constructor(config: GrokConfig) {
    const baseConfig: BaseAgentConfig = {
      sandboxProvider: config.sandboxProvider,
      secrets: config.secrets,
      sandboxId: config.sandboxId,
      workingDirectory: config.workingDirectory,
    };

    super(baseConfig);

    // Validate that provider is xai if specified (Grok only supports xai)
    if (config.provider && config.provider !== "xai") {
      throw new Error("Grok agent only supports 'xai' provider");
    }

    // Store config values
    const envApiKey = process.env.GROK_API_KEY || process.env.XAI_API_KEY;
    const apiKey = config.providerApiKey || envApiKey;
    
    if (!apiKey) {
      throw new Error(
        "Grok agent requires providerApiKey or GROK_API_KEY environment variable"
      );
    }
    
    this.grokApiKey = apiKey;
    this.model = config.model;
    this.baseUrl = config.baseUrl;
  }

  protected getCommandConfig(
    prompt: string,
    mode?: "ask" | "code"
  ): AgentCommandConfig {
    let instruction: string;
    if (mode === "ask") {
      instruction =
        "Research the repository and answer the user's questions. " +
        "Do NOT make any changes to any files in the repository.";
    } else {
      instruction =
        "Do the necessary changes to the codebase based on the users input.\n" +
        "Don't ask any follow up questions.";
    }

    const escapedPrompt = this.escapePrompt(prompt);
    const modelFlag = this.model ? `--model ${this.model}` : "";
    const baseUrlFlag = this.baseUrl ? `--base-url ${this.baseUrl}` : "";

    return {
      command: `echo "${escapedPrompt}" | grok ${modelFlag} ${baseUrlFlag} --prompt "${this.escapePrompt(instruction)}"`,
      errorPrefix: "Grok",
      labelName: "grok",
      labelColor: "00D4AA",
      labelDescription: "Generated by Grok AI agent",
    };
  }

  protected getDefaultTemplate(): string {
    return "vibekit-grok";
  }

  protected getEnvironmentVariables(): Record<string, string> {
    const envVars: Record<string, string> = {
      GROK_API_KEY: this.grokApiKey,
    };

    if (this.baseUrl) {
      envVars.GROK_BASE_URL = this.baseUrl;
    }

    return envVars;
  }

  protected getApiKey(): string {
    return this.grokApiKey;
  }

  protected getAgentType(): AgentType {
    return "grok";
  }

  protected getModelConfig(): ModelConfig {
    return {
      provider: "xai",
      apiKey: this.grokApiKey,
      model: this.model,
      baseUrl: this.baseUrl,
    };
  }

  // Override generateCode to support history in the instruction
  public async generateCode(
    prompt: string,
    mode?: "ask" | "code",
    branch?: string,
    history?: Conversation[],
    callbacks?: GrokStreamCallbacks,
    background?: boolean
  ): Promise<GrokResponse> {
    let instruction: string;
    if (mode === "ask") {
      instruction =
        "Research the repository and answer the user's questions. " +
        "Do NOT make any changes to any files in the repository.";
    } else {
      instruction =
        "Do the necessary changes to the codebase based on the users input.\n" +
        "Don't ask any follow up questions.";
    }

    if (history && history.length > 0) {
      instruction += `\n\nConversation history: ${history
        .map((h) => `${h.role}\n ${this.escapePrompt(h.content)}`)
        .join("\n\n")}`;
    }

    const escapedPrompt = this.escapePrompt(prompt);
    const modelFlag = this.model ? `--model ${this.model}` : "";
    const baseUrlFlag = this.baseUrl ? `--base-url ${this.baseUrl}` : "";

    // Override the command config with history-aware instruction
    const originalGetCommandConfig = this.getCommandConfig.bind(this);
    this.getCommandConfig = (p: string, m?: "ask" | "code") => ({
      ...originalGetCommandConfig(p, m),
      command: `echo "${escapedPrompt}" | grok ${modelFlag} ${baseUrlFlag} --prompt "${this.escapePrompt(instruction)}"`,
    });

    const result = await super.generateCode(
      prompt,
      mode,
      branch,
      history,
      callbacks,
      background
    );

    // Restore original method
    this.getCommandConfig = originalGetCommandConfig;

    return result as GrokResponse;
  }
}